# 获取划分好'train', 'val', 'test'的voc格式数据集中所有的类别及其对应的个数
import os
import xml.etree.ElementTree as ET
from collections import defaultdict

# 定义标签文件夹路径
label_dir = '/home/shuai/Documents/mydataset/my_usv_aug/final_dataset1/labels'

# 定义子文件夹列表
sub_folders = ['train', 'val', 'test']

# 用于存储分类及其数量的字典
class_counts = defaultdict(int)

# 遍历每个子文件夹
for sub_folder in sub_folders:
    sub_folder_path = os.path.join(label_dir, sub_folder)
    
    # 遍历子文件夹中的所有文件
    for file_name in os.listdir(sub_folder_path):
        file_path = os.path.join(sub_folder_path, file_name)
        
        # 解析 XML 文件
        tree = ET.parse(file_path)
        root = tree.getroot()
        
        # 遍历所有的 <object> 标签
        for obj in root.findall('object'):
            # 获取分类名称
            class_name = obj.find('name').text
            # 更新分类数量
            class_counts[class_name] += 1

# 输出分类及其数量统计
print("分类及其数量统计：")
for class_name, count in class_counts.items():
    print(f"{class_name}: {count}")

# 提取所有分类名称并转换为列表
class_list = list(class_counts.keys())

# 输出所有分类列表
print("\n所有分类列表：")
print(class_list)